# @Author: Zhixuan.Wang
# @IDE: PyCharm
# @Project: multimodal
# @File: DataLoader.py
# @Time: 2025/11/9 15:40
# @Description: 训练集测试机数据加载

import torch
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from MyDataSet import MyImgDataset

def collate_fn(batch):
    """用于处理不同大小目标的 batch"""
    images, labels, bboxes = zip(*batch)
    return list(images), list(labels), list(bboxes)


def get_dataloaders(
        img_path: str = r"D:\LungCancer\VOCdevkit\VOC2012\EffectiveImage",
        xml_path: str = r"D:\LungCancer\VOCdevkit\VOC2012\NEW-Annotations",
        batch_size: int = 4,
        train_ratio: float = 0.7,
        num_workers: int = 0  # ⚠️ Windows 建议默认设为 0，避免多进程问题
):
    """
    构建训练集与验证集的 DataLoader
    Args:
        img_path (str): 图像目录
        xml_path (str): 标注 XML 目录
        batch_size (int): 批大小
        train_ratio (float): 训练集比例
        num_workers (int): 加载线程数（Windows建议0）
    Returns:
        train_data, val_data
    """

    # ===== 数据增强 / 预处理 =====
    transform = transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 三通道
    ])

    # ===== 构建数据集 =====
    dataset = MyImgDataset(
        img_path=img_path,
        xml_path=xml_path,
        transform=transform
    )

    # ===== 划分训练集 / 验证集 =====
    train_size = int(train_ratio * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

    # ===== 构建 DataLoader =====
    train_data = DataLoader(train_dataset, batch_size=batch_size,
                            shuffle=True, num_workers=num_workers,
                            pin_memory=True, collate_fn=collate_fn)
    val_data = DataLoader(val_dataset, batch_size=batch_size,
                          shuffle=False, num_workers=num_workers,
                          pin_memory=True, collate_fn=collate_fn)

    return train_data, val_data


# if __name__ == "__main__":
#     # 测试
#     train_loader, val_loader = get_dataloaders(
#         img_path=r"D:\LungCancer\VOCdevkit\VOC2012\EffectiveImage",
#         xml_path=r"D:\LungCancer\VOCdevkit\VOC2012\NEW-Annotations",
#         batch_size=4
#     )
#
#     print(f"Train batches: {len(train_loader)} | Val batches: {len(val_loader)}")
#     for imgs, labels, boxes in train_loader:
#         print("✅ Loaded batch:", imgs[0].shape, labels[0], boxes[0])
#         break
